#ifndef SORT3D_HPP
#define SORT3D_HPP

#include <eigen3/Eigen/Dense>
#include <vector>
#include <cmath>
#include <algorithm>
#include <numeric>
#include "tracker/lapjv.hpp" // linear assignment problem solver

// This class represents the internal state of individual tracked objects in 3D
class KalmanTracker
{
public:
  static int count;
  int id, hits, hit_streak;
  double time_since_update;

  // Initializes a tracker using initial position
  explicit KalmanTracker(const Eigen::Vector3d &xyz) {
    // State transition matrix I + F*dt
    I = Eigen::MatrixXd::Identity(6, 6);
    F = Eigen::MatrixXd::Zero(6, 6);
    F(0, 3) = F(1, 4) = F(2, 5) = 1.0;

    // Measurement matrix
    H = Eigen::MatrixXd::Zero(3, 6);
    H(0, 0) = H(1, 1) = H(2, 2) = 1.0;

    // Initial state uncertainty
    P = Eigen::MatrixXd::Identity(6, 6) * 10.0;
    P.block<3, 3>(3, 3) *= 1000.0; // Velocity uncertainty

    // Process uncertainty
    Q = Eigen::MatrixXd::Identity(6, 6) * 0.01;

    // Measurement uncertainty
    R = Eigen::MatrixXd::Identity(3, 3) * 0.5;

    // Initial state zero
    x = Eigen::VectorXd(6).setZero();
    x.head<3>() = xyz;

    time_since_update = 0.0;
    id = count++;
    hits = 0;
    hit_streak = 0;
  }

  // Updates the state vector with observed 3D position
  void update(const Eigen::Vector3d &xyz) {
    // y = z - Hx
    // error (residual) between measurement and prediction
    Eigen::MatrixXd y = xyz - H * x;

    Eigen::MatrixXd PHT = P * H.transpose();
    // S = HPH' + R
    // project system uncertainty into measurement space
    Eigen::MatrixXd S = H * PHT + R;

    // K = PH'inv(S)
    // map system uncertainty into kalman gain
    Eigen::MatrixXd K = PHT * S.inverse();

    // x = x + Ky
    // predict new x with residual scaled by the kalman gain
    x = x + K * y;

    // P = (I-KH)P(I-KH)' + KRK'
    // This is more numerically stable
    // and works for non-optimal K vs the equation
    // P = (I-KH)P usually seen in the literature.
    Eigen::MatrixXd I_KH = Eigen::MatrixXd::Identity(6, 6) - K * H;
    P = I_KH * P * I_KH.transpose() + K * R * K.transpose();
    // P = I_KH * P;

    time_since_update = 0.0;
    hits++;
    hit_streak++;
  }

  // Advances the state vector and returns the predicted 3D position estimate
  Eigen::Vector3d predict(double dt = 1.0) {
    Eigen::MatrixXd _F = I + F * dt;
    // x = Fx + Bu
    x = _F * x;
    // P = FPF' + Q
    P = _F * P * _F.transpose() + Q;

    if (time_since_update > 0) hit_streak = 0;
    time_since_update += dt;
    return x.head<3>();
  }

  // Returns the current 3D position estimate.
  Eigen::Vector3d getState() const { return x.head<3>(); }

private:
  Eigen::MatrixXd I, F, H, P, Q, R;
  Eigen::VectorXd x;
};

int KalmanTracker::count = 0;

class Sort3D
{
public:
  Sort3D() = default;

  /**
   * @param max_age   maximum time to keep the tracker alive w/o (second)
   *                  associated detections
   * @param min_hits  minimum number of associated detections before track is
   *                  considered valid
   * @param dist_thresh distance threshold for data association
   */
  void init(double max_age, int min_hits, double dist_thresh) {
    max_age_ = max_age;
    min_hits_ = min_hits;
    dist_thresh_ = dist_thresh;
  }

  std::vector<Eigen::Vector4d> update(const std::vector<Eigen::Vector3d> &dets,
                                      double dt = 1.0) {
    frame_count_++;

    std::vector<Eigen::Vector3d> trks;

    // Get predicted locations from existing trackers.
    trackers.erase(std::remove_if(trackers.begin(),
                                  trackers.end(),
                                  [&trks, dt](KalmanTracker &tracker) {
                                    Eigen::Vector3d pos = tracker.predict(dt);
                                    if (pos.hasNaN())
                                      return true; // remove tracker if NaN
                                    else {
                                      trks.emplace_back(pos);
                                      return false;
                                    }
                                  }),
                   trackers.end());

    std::vector<Eigen::Vector2i> matches;
    std::vector<int> unmatched_detections;

    associate_detections_to_trackers(dets, trks, matches, unmatched_detections);

    // update matched trackers with assigned detections
    for (auto mat : matches) {
      trackers[mat(1)].update(dets[mat(0)]);
    }

    // create and initialise new trackers for unmatched detections
    for (auto i : unmatched_detections) {
      trackers.emplace_back(dets[i]);
    }

    // delete lost trackers
    double _max_age = max_age_;
    trackers.erase(std::remove_if(trackers.begin(),
                                  trackers.end(),
                                  [_max_age](const KalmanTracker &tracker) {
                                    return tracker.time_since_update > _max_age;
                                  }),
                   trackers.end());

    // return the set of tracks
    std::vector<Eigen::Vector4d> id_poses;
    id_poses.reserve(trackers.size());

    for (const auto &tracker : trackers) {
      if (tracker.time_since_update == 0.0 &&
          (tracker.hit_streak >= min_hits_ || frame_count_ <= min_hits_)) {
        Eigen::VectorXd row(4);
        row << tracker.id, tracker.getState();
        id_poses.emplace_back(row);
      }
    }

    return id_poses;
  }

private:
  double max_age_{2.5};
  int min_hits_{0}, frame_count_{0};
  double dist_thresh_{2.0};
  std::vector<KalmanTracker> trackers;

  void associate_detections_to_trackers(
      const std::vector<Eigen::Vector3d> &detections,
      const std::vector<Eigen::Vector3d> &trks,
      std::vector<Eigen::Vector2i> &matches,
      std::vector<int> &unmatched_detections) {
    matches.clear();
    unmatched_detections.clear();

    if (detections.empty()) return;

    if (trks.empty()) {
      unmatched_detections.resize(detections.size());
      std::iota(unmatched_detections.begin(), unmatched_detections.end(), 0);
      return;
    }

    // 这样更快，详细对比见 test/test1.cpp
    Eigen::MatrixXd distance_matrix(detections.size(), trks.size());
    for (size_t i = 0; i < detections.size(); ++i) {
      for (size_t j = 0; j < trks.size(); ++j) {
        distance_matrix(i, j) = (detections[i] - trks[j]).norm();
      }
    }

    // use linear assignment to find matches(lapjv)
    int dim = std::max(distance_matrix.rows(), distance_matrix.cols());
    std::vector<int> rowsol(dim), colsol(dim);
    std::vector<double> u(dim), v(dim);
    double *assign_cost = fillMatrix(distance_matrix);
    lap<false>(
        dim, assign_cost, rowsol.data(), colsol.data(), u.data(), v.data());
    delete[] assign_cost;

    for (int i = 0; i < distance_matrix.rows(); i++) {
      if (rowsol[i] < distance_matrix.cols() &&
          distance_matrix(i, rowsol[i]) < dist_thresh_) {
        matches.emplace_back(Eigen::Vector2i(i, rowsol[i]));
      } else {
        unmatched_detections.emplace_back(i);
      }
    }
  }

  double *fillMatrix(const Eigen::MatrixXd &costMatrix) {
    int rows = costMatrix.rows();
    int cols = costMatrix.cols();
    int dim = std::max(rows, cols);
    auto *assign_cost = new double[dim * dim];

    for (int i = 0; i < dim; ++i) {
      for (int j = 0; j < dim; ++j) {
        if (i < rows && j < cols) {
          assign_cost[i * dim + j] = costMatrix(i, j);
        } else {
          assign_cost[i * dim + j] = 0.0;
        }
      }
    }
    return assign_cost;
  }
};

#endif // SORT3D_HPP
